package burp;

import burp.Utils.JsonUtils;
import burp.Utils.LoggerUtils;
import burp.Utils.RequestUtils;

import javax.swing.*;
import java.util.*;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@SuppressWarnings("unused")
public class BurpExtender implements IBurpExtender, IContextMenuFactory {
    private static BurpExtender instance;
    private IBurpExtenderCallbacks callbacks;
    private IExtensionHelpers helpers;
    private ParamNullifierHistory historyTab;

    // 配置项（由 Settings 对话框修改）
    public static final Set<String> TARGET_HEADERS = new HashSet<>();
    public static final Set<String> TARGET_COOKIES = new HashSet<>();
    public static final Set<String> TARGET_URL_PARAMS = new HashSet<>();
    public static final Set<String> TARGET_BODY_PARAMS = new HashSet<>();
    public static boolean CHECK_RESPONSE_CODE = true;
    public static final Set<String> CUSTOM_RESPONSE_INDICATORS = new HashSet<>();

    // 线程池
    static final ExecutorService EXECUTOR = Executors.newFixedThreadPool(5); // 可根据需要调整线程数

    // 原始请求集合
    private final List<IHttpRequestResponse> requestResponses = new ArrayList<>();

    // 获取单例
    public static BurpExtender getInstance() {
        if (instance == null) {
            throw new IllegalStateException("BurpExtender not initialized");
        }
        return instance;
    }

    // 获取 IExtensionHelpers
    public IExtensionHelpers getHelpers() {
        return helpers;
    }

    // 获取 IBurpExtenderCallbacks
    public IBurpExtenderCallbacks getCallbacks() {
        return callbacks;
    }

    // 插件入口
    @Override
    public void registerExtenderCallbacks(IBurpExtenderCallbacks callbacks) {
        instance = this;
        this.callbacks = callbacks;
        this.helpers = callbacks.getHelpers();

        // 初始化历史标签页
        this.historyTab = new ParamNullifierHistory(requestResponses);
        callbacks.addSuiteTab(historyTab);

        callbacks.setExtensionName("ParamNullifier");
        callbacks.registerContextMenuFactory(this);
        LoggerUtils.info("Plugin loaded");
    }

    // 创建右键菜单
    @Override
    public List<JMenuItem> createMenuItems(IContextMenuInvocation invocation) {
        int context = invocation.getInvocationContext();
        if (context == IContextMenuInvocation.CONTEXT_MESSAGE_EDITOR_REQUEST ||
                context == IContextMenuInvocation.CONTEXT_MESSAGE_VIEWER_REQUEST ||
                context == IContextMenuInvocation.CONTEXT_PROXY_HISTORY ||
                context == IContextMenuInvocation.CONTEXT_TARGET_SITE_MAP_TABLE) {

            List<JMenuItem> menus = new ArrayList<>();
            JMenuItem sendToHistoryItem = getJMenuItem(invocation);
            menus.add(sendToHistoryItem);
            return menus;
        }
        return null;
    }

    // 将原始请求加入历史记录
    private JMenuItem getJMenuItem(IContextMenuInvocation invocation) {
        JMenuItem sendToHistoryItem = new JMenuItem("Send to ParamNullifierHistory");
        sendToHistoryItem.addActionListener(e -> {
            final IHttpRequestResponse[] messages = invocation.getSelectedMessages();
            if (messages == null || messages.length == 0) return;

            for (IHttpRequestResponse message : messages) {
                IHttpRequestResponse savedMessage = callbacks.saveBuffersToTempFiles(message);
                requestResponses.add(savedMessage); // 存储原始请求
                historyTab.addOriginalEntry(savedMessage); // 显示在原始请求表格中
            }
            LoggerUtils.info("Request added to history");
        });
        return sendToHistoryItem;
    }

    public void sendNullifiedRequest(
            IHttpRequestResponse originalMessage,
            String paramName,
            boolean isUrlParam,
            boolean isHeader,
            boolean isCookie,
            CountDownLatch latch) {
        if (originalMessage == null || originalMessage.getRequest() == null) {
            LoggerUtils.error("Invalid original message");
            latch.countDown();
            return;
        }

        IExtensionHelpers helpers = this.helpers;
        IRequestInfo requestInfo = helpers.analyzeRequest(originalMessage.getHttpService(), originalMessage.getRequest());
        if (requestInfo == null) {
            LoggerUtils.error("Invalid request info");
            latch.countDown();
            return;
        }
        IHttpService httpService = originalMessage.getHttpService();
        if (httpService == null) {
            LoggerUtils.error("Invalid http service");
            latch.countDown();
            return;
        }

        // 修改所有参数和构造新请求
        byte[] newRequestBytes = RequestUtils.nullifyParameter(helpers, requestInfo, originalMessage.getRequest(), paramName, isUrlParam, isHeader, isCookie);
        if (newRequestBytes == null) {
            LoggerUtils.error("Failed to nullify parameter");
            latch.countDown();
            return;
        }

        // 发送新请求
        IHttpRequestResponse nullifiedMessage = callbacks.makeHttpRequest(httpService, newRequestBytes);
        if (nullifiedMessage == null) {
            LoggerUtils.error("makeHttpRequest returned null");
            latch.countDown();
            return;
        }

        // 保存并更新 UI
        IHttpRequestResponse savedNullifiedMessage = callbacks.saveBuffersToTempFiles(nullifiedMessage);
        if (savedNullifiedMessage != null && savedNullifiedMessage.getRequest() != null) {
            SwingUtilities.invokeLater(() -> historyTab.addNullifiedEntry(savedNullifiedMessage));
        }

        // 实际请求完成后再 countDown
        latch.countDown();
    }

    /**
     * 测试所有参数的置空请求
     */
    public void testAllNullifiedParams(IHttpRequestResponse selectedMessage, Runnable onCompletion) {
        if (selectedMessage == null || selectedMessage.getRequest() == null) {
            LoggerUtils.warn("Invalid request");
            onCompletion.run();
            return;
        }

        // 深拷贝原始请求
        IHttpRequestResponse savedMessage = callbacks.saveBuffersToTempFiles(selectedMessage);
        if (savedMessage == null) {
            LoggerUtils.error("Failed to save message");
            onCompletion.run();
            return;
        }

        // 获取请求信息
        IRequestInfo requestInfo = helpers.analyzeRequest(savedMessage.getHttpService(), savedMessage.getRequest());
        List<String> headers = new ArrayList<>(requestInfo.getHeaders());
        String contentType = RequestUtils.getHeaderValue(headers, "Content-Type");
        boolean isJson = contentType != null && contentType.contains("application/json");
        boolean isFormPost = contentType != null && contentType.contains("application/x-www-form-urlencoded");

        // 创建线程池
        Set<String> processedParams = Collections.synchronizedSet(new HashSet<>());
        int[] paramCount = {0};

        // 计算要处理的参数总数
        synchronized (processedParams) {
            if (!TARGET_HEADERS.isEmpty()) paramCount[0] += TARGET_HEADERS.size();
            if (!TARGET_COOKIES.isEmpty()) paramCount[0] += TARGET_COOKIES.size();

            if (TARGET_URL_PARAMS.isEmpty()) {
                for (IParameter param : requestInfo.getParameters()) {
                    if (param.getType() == IParameter.PARAM_URL) paramCount[0]++;
                }
            } else {
                paramCount[0] += TARGET_URL_PARAMS.size();
            }

            if (isFormPost) {
                if (TARGET_BODY_PARAMS.isEmpty()) {
                    for (IParameter param : requestInfo.getParameters()) {
                        if (param.getType() == IParameter.PARAM_BODY) paramCount[0]++;
                    }
                } else {
                    paramCount[0] += TARGET_BODY_PARAMS.size();
                }
            } else if (isJson) {
                int jsonBodyOffset = requestInfo.getBodyOffset();
                byte[] jsonBodyBytes = Arrays.copyOfRange(savedMessage.getRequest(), jsonBodyOffset, savedMessage.getRequest().length);
                Map<String, Object> jsonMap = JsonUtils.parseJsonBody(helpers.bytesToString(jsonBodyBytes));
                if (TARGET_BODY_PARAMS.isEmpty()) {
                    paramCount[0] += jsonMap.size();
                } else {
                    paramCount[0] += TARGET_BODY_PARAMS.size();
                }
            }
        }

        // 线程池中任务的数量
        CountDownLatch latch = new CountDownLatch(paramCount[0]);

        // 提交到线程池执行
        EXECUTOR.submit(() -> {
            try {
                long startTime = System.currentTimeMillis();

                // 处理 TARGET_HEADERS
                for (String headerName : TARGET_HEADERS) {
                    boolean found = false;
                    for (String header : headers) {
                        if (header.startsWith(headerName + ":")) {
                            found = true;
                            break;
                        }
                    }
                    if (found) {
                        sendNullifiedRequest(savedMessage, headerName, false, true, false, latch);
                        processedParams.add(headerName);
                    } else {
                        LoggerUtils.info("Skipped header not found: " + headerName);
                        latch.countDown();
                    }
                }

                // 处理 TARGET_COOKIES
                for (String cookieName : TARGET_COOKIES) {
                    boolean found = false;
                    for (String header : requestInfo.getHeaders()) {
                        if (header.startsWith("Cookie:") && header.contains(cookieName + "=")) {
                            found = true;
                            break;
                        }
                    }
                    if (found) {
                        sendNullifiedRequest(savedMessage, cookieName, false, false, true, latch);
                        processedParams.add(cookieName);
                    } else {
                        LoggerUtils.info("Skipped cookie not found: " + cookieName);
                        latch.countDown();
                    }
                }

                // 自动识别 URL 参数
                if (TARGET_URL_PARAMS.isEmpty()) {
                    for (IParameter param : requestInfo.getParameters()) {
                        if (param.getType() == IParameter.PARAM_URL) {
                            sendNullifiedRequest(savedMessage, param.getName(), true, false, false, latch);
                            processedParams.add(param.getName());
                        }
                    }
                } else {
                    for (String urlParam : TARGET_URL_PARAMS) {
                        for (IParameter param : requestInfo.getParameters()) {
                            if (param.getType() == IParameter.PARAM_URL && param.getName().equals(urlParam)) {
                                sendNullifiedRequest(savedMessage, urlParam, true, false, false, latch);
                                processedParams.add(urlParam);
                                break;
                            }
                        }
                    }
                }

                // 处理 Body 参数
                if (isFormPost) {
                    if (TARGET_BODY_PARAMS.isEmpty()) {
                        for (IParameter param : requestInfo.getParameters()) {
                            if (param.getType() == IParameter.PARAM_BODY) {
                                sendNullifiedRequest(savedMessage, param.getName(), false, false, false, latch);
                                processedParams.add(param.getName());
                            }
                        }
                    } else {
                        for (String bodyParam : TARGET_BODY_PARAMS) {
                            for (IParameter param : requestInfo.getParameters()) {
                                if (param.getType() == IParameter.PARAM_BODY && param.getName().equals(bodyParam)) {
                                    sendNullifiedRequest(savedMessage, bodyParam, false, false, false, latch);
                                    processedParams.add(bodyParam);
                                    break;
                                }
                            }
                        }
                    }
                } else if (isJson) {
                    int jsonBodyOffset = requestInfo.getBodyOffset();
                    byte[] jsonBodyBytes = Arrays.copyOfRange(savedMessage.getRequest(), jsonBodyOffset, savedMessage.getRequest().length);
                    Map<String, Object> jsonMap = JsonUtils.parseJsonBody(helpers.bytesToString(jsonBodyBytes));

                    if (TARGET_BODY_PARAMS.isEmpty()) {
                        for (String key : jsonMap.keySet()) {
                            sendNullifiedRequest(savedMessage, key, false, false, false, latch);
                            processedParams.add(key);
                        }
                    } else {
                        for (String bodyParam : TARGET_BODY_PARAMS) {
                            if (jsonMap.containsKey(bodyParam)) {
                                sendNullifiedRequest(savedMessage, bodyParam, false, false, false, latch);
                                processedParams.add(bodyParam);
                            } else {
                                latch.countDown();
                            }
                        }
                    }
                }

                // 等待所有任务完成
                latch.await();

                // 输出统计信息
                long responseTime = System.currentTimeMillis() - startTime;
                SwingUtilities.invokeLater(() -> {
                    if (processedParams.isEmpty()) {
                        LoggerUtils.info("No matching parameters to nullify in this request");
                    } else {
                        LoggerUtils.info("Processed " + processedParams.size() + " parameters in " + responseTime + "ms");
                    }
                });

                // 所有任务完成，回调通知
                onCompletion.run();

            } catch (Exception e) {
                LoggerUtils.error("Error processing request", e);
                onCompletion.run();
            }
        });
    }


    /**
     * 清空原始请求
     */
    public void clearOriginalRequests() {
        requestResponses.clear();
    }
}
